#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File     : Models.py
# @Time     : 2024/8/14 15:07


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from torchtext.legacy.data import Field, Dataset, BucketIterator
from torchtext.legacy.datasets import TranslationDataset

from Layers import *


def get_pad_mask(seq, pad_idx):
    """
    用于生成一个填充掩码，通过指示哪些位置是实际单词，哪些是填充标记，帮助模型在计算注意力时忽略填充部分。
    return: (batch_size, 1, seq_len)
    """
    return (seq != pad_idx).unsqueeze(-2)


def get_subsequent_mask(seq):
    '''
    生成一个下三角矩阵，确保模型在生成下一个单词时只考虑当前和之前的单词，而不查看未来的单词
    具体步骤：
        1、创建一个形状为 (1, len_s, len_s) 的全 1 张量
        2、torch.triu: 用于返回一个上三角矩阵，其中 diagonal=1 表示对角线以上的部分保留 1，下面的部分设为 0。
        3、1 - torch.triu: 对上三角矩阵进行取反操作，将 1 变为 0，将 0 变为 1，变为下三角矩阵
        4、.bool(): 张量转换为布尔型，生成掩码矩阵。
    '''
    sz_b, len_s = seq.size()
    subsequent_mask = (1 - torch.triu(
        torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool()
    return subsequent_mask


class PositionalEncoding(nn.Module):
    """位置编码"""

    def __init__(self, d_hid, n_position=200):
        super(PositionalEncoding, self).__init__()
        # register_buffer 注册的张量不会在优化时被更新（不会在 optimizer.step() 时被调整），但它仍然是模型的一部分，
        # 并且在模型保存和加载时会被保留。
        #   1. pos_table 会在模型的保存和加载中被包括。
        #   2. 但是 pos_table 不会被视为可训练参数，不会被优化器更新。
        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        """正弦位置编码表"""

        def get_position_angle_vec(position):
            """
            位置向量的第多少维                        i               2i             2i+1
            (0 2 4等偶数维用sin函数计算)
                    0                             i=0//2=0         0
                    1                             i=1//2=0                         1
                    2                             i=2//2=1         2
                    3                             i=3//2=1                         3
                    4                             i=4//2=2         4
                    5                             i=5//2=2                         5

                    510                           i=510//2=255     510
                    511                           i=511//2=255                     511
            2i   是指向量维度中的偶数维，即第0维、第2维、第4维...，第510维，用sin函数计算
            2i+1 是指向量维度中的奇数维，即第1维、第3维、第5维..，第511维，用cos函数计算
            """
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.tensor(sinusoid_table, dtype=torch.float32).unsqueeze(0)  # (1,N,d)

    def forward(self, x):
        # x(B,N,D)
        #   1、:x.size(1)这会从位置编码表中选择与输入序列长度相匹配[N]的部分,如果输入序列长度小于最大位置，只会使用需要的部分
        #   2、detach()是为了确保位置编码不参与梯度计算。这是因为位置编码是预先计算的固定值，不需要在训练过程中更新。
        # 将输入与embedding相加
        return x + self.pos_table[:, :x.size(1)].clone().detach()


class Encoder(nn.Module):
    def __init__(self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
                 d_model, d_inner, pad_idx, dropout=0.1, n_position=200, scale_emb=False):
        """
        d_word_vec:词向量的维度
        """
        super(Encoder, self).__init__()
        self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.droput = nn.Dropout(dropout)
        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)
        ])
        self.scale_emb = scale_emb
        self.d_model = d_model
        self.layer_norm = nn.LayerNorm(d_model, 1e-6)

    def forward(self, src_seq, src_mask, return_attns=False):
        enc_slf_attn_list = []
        enc_ouput = self.src_word_emb(src_seq)

        if self.scale_emb:
            enc_ouput *= self.d_model ** 0.5

        enc_ouput = self.droput(self.position_enc(enc_ouput))
        enc_ouput = self.layer_norm(enc_ouput)

        for enc_layer in self.layer_stack:
            enc_ouput, enc_slf_attn = enc_layer(enc_ouput, slf_attn_mask=src_mask)
            enc_slf_attn_list += [enc_slf_attn] if return_attns else []

        if return_attns:
            return enc_ouput, enc_slf_attn_list
        return enc_ouput


class Decoder(nn.Module):
    def __init__(
            self, n_trg_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, n_position=200, dropout=0.1, scale_emb=False):
        super(Decoder, self).__init__()
        self.trg_word_emb = nn.Embedding(n_trg_vocab, d_word_vec, padding_idx=pad_idx)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(dropout)
        self.layer_stack = nn.ModuleList([
            DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)
        ])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, trg_seq, trg_mask, enc_output, src_mask, return_attns=False):
        dec_slf_attn_list, dec_enc_attn_list = [], []

        dec_output = self.trg_word_emb(trg_seq)

        if self.scale_emb:
            # 对目标语言的词嵌入进行缩放
            dec_output *= self.d_model ** 0.5

        dec_output = self.dropout(self.position_enc(dec_output))
        dec_output = self.layer_norm(dec_output)

        for dec_layer in self.layer_stack:
            dec_output, dec_slf_attn, dec_enc_attn = dec_layer(
                dec_output, enc_output, slf_attn_mask=trg_mask, dec_enc_attn_mask=src_mask)
            dec_slf_attn_list += [dec_slf_attn] if return_attns else []
            dec_enc_attn_list += [dec_enc_attn] if return_attns else []

        if return_attns:
            return dec_output, dec_slf_attn_list, dec_enc_attn_list
        return dec_output


class Transformer(nn.Module):
    def __init__(
            self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx,
            d_word_vec=512, d_model=512, d_inner=2048,
            n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200,
            trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True,
            scale_emb_or_prj='prj'):
        """
        n_src_vocab: 源语言（例如英语）的词汇表大小。
        n_trg_vocab: 目标语言（例如中文）的词汇表大小。

        trg_emb_prj_weight_sharing: 是否在目标词嵌入层和输出层之间共享权重，默认值为 True。
        emb_src_trg_weight_sharing: 是否在源词嵌入层和目标词嵌入层之间共享权重，默认值为 True。
        scale_emb_or_prj: 这是一个字符串参数，决定了模型中某些层的输出是否要进行缩放处理。有三个选项：
            'emb': 对嵌入层的输出乘以 √d_model。
            'prj': 对线性投影层的输出乘以 √d_model^-1。
            'none': 不进行缩放操作。
        """
        super(Transformer, self).__init__()
        self.src_pad_idx, self.trg_pad_idx = src_pad_idx, trg_pad_idx

        assert scale_emb_or_prj in ['emb', 'prj', 'none']
        scale_emb = (scale_emb_or_prj == 'emb') if trg_emb_prj_weight_sharing else False
        print(f'scale_emb:{scale_emb}')  # False
        self.scale_prj = (scale_emb_or_prj == 'prj') if trg_emb_prj_weight_sharing else False
        self.d_model = d_model

        self.encoder = Encoder(
            n_src_vocab=n_src_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            pad_idx=src_pad_idx, dropout=dropout, scale_emb=scale_emb)

        self.decoder = Decoder(
            n_trg_vocab=n_trg_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            pad_idx=trg_pad_idx, dropout=dropout, scale_emb=scale_emb)

        # 将解码器输出映射到目标语言词汇表的线性层。输出的维度为目标语言词汇表大小
        self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

        # 对神经网络模型中的参数进行初始化，‌具体是对模型中的所有参数进行Xavier均匀初始化。‌
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

        assert d_model == d_word_vec, \
            'To facilitate the residual connections, \
                     the dimensions of all module outputs shall be the same.'

        if trg_emb_prj_weight_sharing:
            # 在目标词嵌入层和输出层之间共享权重
            self.trg_word_prj.weight = self.decoder.trg_word_emb.weight

        if emb_src_trg_weight_sharing:
            # 在源词嵌入层和目标词嵌入层之间共享权重
            self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight

    def forward(self, src_seq, trg_seq):
        # 用于产生Encoder的Mask，它是一列Bool值，负责把标点mask掉
        src_mask = get_pad_mask(src_seq, self.src_pad_idx)

        # 1.(batch_size, 1, seq_len)，遮蔽填充标记的位置。
        # 2.(1, seq_len, seq_len)，遮蔽未来信息的位置。
        # 3.(batch_size, seq_len, seq_len) &:都为 True 的位置在结果中保持 True
        trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)

        enc_output, _ = self.encoder(src_seq, src_mask)
        dec_output, _ = self.decoder(trg_seq, trg_mask, enc_output, src_mask)

        # (batch_size, trg_seq_len, d_model) --> (batch_size, trg_seq_len, n_trg_vocab)
        seq_logit = self.trg_word_prj(dec_output)

        if self.scale_prj:
            # 'prj': 对线性投影层的输出乘以 √d_model ^ -1。
            seq_logit *= self.d_model ** -0.5

        # (batch_size, trg_seq_len, n_trg_vocab) --> (batch_size * n_trg_vocab, trg_seq_len)
        # view(-1, xx) 表示自动计算-1所在的维度
        return seq_logit.view(-1, seq_logit.size(-2))
